package org.dousb.doupicturebackend.manager.sharding;


import com.baomidou.mybatisplus.extension.toolkit.SqlRunner;
import lombok.extern.slf4j.Slf4j;
import org.apache.shardingsphere.driver.jdbc.core.connection.ShardingSphereConnection;
import org.apache.shardingsphere.infra.metadata.database.rule.ShardingSphereRuleMetaData;
import org.apache.shardingsphere.mode.manager.ContextManager;
import org.apache.shardingsphere.sharding.api.config.ShardingRuleConfiguration;
import org.apache.shardingsphere.sharding.api.config.rule.ShardingTableRuleConfiguration;
import org.apache.shardingsphere.sharding.rule.ShardingRule;
import org.dousb.doupicturebackend.model.entity.Space;
import org.dousb.doupicturebackend.model.enums.SpaceLevelEnum;
import org.dousb.doupicturebackend.model.enums.SpaceTypeEnum;
import org.dousb.doupicturebackend.service.SpaceService;
import org.springframework.stereotype.Component;

import javax.annotation.PostConstruct;
import javax.annotation.Resource;
import javax.sql.DataSource;
import java.sql.SQLException;
import java.util.Collections;
import java.util.List;
import java.util.Optional;
import java.util.Set;
import java.util.stream.Collectors;

/**
 * 分表管理器
 * <p>
 * 该类负责管理ShardingSphere的分表配置，实现动态分表功能。
 * 主要功能包括：
 * 1. 初始化分表配置
 * 2. 获取所有动态表名
 * 3. 更新动态表名配置
 * 4. 创建空间分表
 */
//@Component
@Slf4j
public class DynamicShardingManager {

    // 逻辑表名常量
    private static final String LOGIC_TABLE_NAME = "picture";
    // 数据库名常量
    private static final String DATABASE_NAME = "logic_db";
    // 数据源注入
    @Resource
    private DataSource dataSource;
    // 空间服务注入
    @Resource
    private SpaceService spaceService;

    /**
     * 初始化方法
     * <p>
     * 在类实例化后自动执行，用于初始化动态分表配置
     */
    @PostConstruct
    public void initialize() {
        log.info("初始化动态分表配置");
        updateShardingTableNodes();
    }

    /**
     * 获取所有的动态表名  包括初始表picture和分表picture_spaceId
     *
     * @return
     */
    private Set<String> fetchAllPictureTableName() {
        //为了测试方便 直接将所有的团队空间分表（实际上线改为 仅旗舰空间）
        Set<Long> spaceIds = spaceService.lambdaQuery()
                .eq(Space::getSpaceType, SpaceTypeEnum.TEAM.getValue())
                .list()
                .stream()
                .map(Space::getId)
                .collect(Collectors.toSet());
        Set<String> tableNames = spaceIds.stream()
                .map(spaceId -> LOGIC_TABLE_NAME + "_" + spaceId)
                .collect(Collectors.toSet());

        //添加初始逻辑表
        tableNames.add(LOGIC_TABLE_NAME);

        return tableNames;
    }

    /**
     * 更新动态表名配置
     * 该方法用于动态更新ShardingSphere中的分表规则配置，根据获取到的表名集合更新实际数据节点配置
     *
     * @return 无返回值
     */
    private void updateShardingTableNodes() {
        Set<String> tableNames = fetchAllPictureTableName();
        String newActualDataNodes = tableNames.stream()
                .map(tableName -> "dou_picture." + tableName)
                .collect(Collectors.joining(","));
        log.info("动态分表配置：{}", newActualDataNodes);

        ContextManager contextManager = getContextManager();
        ShardingSphereRuleMetaData ruleMetaData = contextManager.getMetaDataContexts()
                .getMetaData()
                .getDatabases()
                .get(DATABASE_NAME)
                .getRuleMetaData();

        Optional<ShardingRule> shardingRule = ruleMetaData.findSingleRule(ShardingRule.class);
        if (shardingRule.isPresent()) {
            ShardingRuleConfiguration ruleConfig = (ShardingRuleConfiguration) shardingRule.get().getConfiguration();
            List<ShardingTableRuleConfiguration> updatedRules = ruleConfig.getTables()
                    .stream()
                    .map(oldTableRule -> {
                        if (LOGIC_TABLE_NAME.equals(oldTableRule.getLogicTable())) {
                            ShardingTableRuleConfiguration newTableRuleConfig = new ShardingTableRuleConfiguration(LOGIC_TABLE_NAME, "");
                            newTableRuleConfig.setDatabaseShardingStrategy(oldTableRule.getDatabaseShardingStrategy());
                            newTableRuleConfig.setTableShardingStrategy(oldTableRule.getTableShardingStrategy());
                            newTableRuleConfig.setKeyGenerateStrategy(oldTableRule.getKeyGenerateStrategy());
                            newTableRuleConfig.setAuditStrategy(oldTableRule.getAuditStrategy());
                            return newTableRuleConfig;
                        }
                        return oldTableRule;
                    })
                    .collect(Collectors.toList());
            ruleConfig.setTables(updatedRules);
            contextManager.alterRuleConfiguration(DATABASE_NAME, Collections.singleton(ruleConfig));
            contextManager.reloadDatabase(DATABASE_NAME);
            log.info("动态分表规则更新成功！");
        } else {
            log.error("未找到 ShardingSphere 的分片规则配置，动态分表更新失败。");
        }

    }

    /**
     * 获取 ShardingSphere ContextManager
     */
    private ContextManager getContextManager() {
        try (ShardingSphereConnection connection = dataSource.getConnection().unwrap(ShardingSphereConnection.class)) {
            return connection.getContextManager();
        } catch (SQLException e) {
            throw new RuntimeException("获取 ShardingSphere ContextManager 失败", e);
        }
    }

    /**
     * 创建空间图片表
     *
     * @param space 空间对象
     */
    public void createSpacePictureTable(Space space) {
        //动态创建分表
        //仅为旗舰版本的团队空间进行分表
        if (space.getSpaceLevel().equals(SpaceLevelEnum.FLAGSHIP.getValue()) && space.getSpaceType().equals(SpaceTypeEnum.TEAM.getValue())) {
            Long spaceId = space.getId();
            String tableName = "picture_" + spaceId;
            //创建新表
            String createTableSql = "CREATE TABLE " + tableName + " LIKE " + LOGIC_TABLE_NAME;
            log.info("准备执行创建分表SQL: {}", createTableSql);
            try {
                SqlRunner.db().update(createTableSql);
                //更新分表
                updateShardingTableNodes();
            } catch (Exception e) {
                e.printStackTrace();
                log.error("创建空间分表失败：{}", spaceId);
            }
        }
    }


}
